import timeit
import numpy as np
import av

aligned = av.VideoFrame(1920, 1080, format="rgb24")
padded = av.VideoFrame(1080, 1920, format="rgb24")
aligned_yuv = av.VideoFrame(1920, 1080, format="yuv420p")
padded_yuv = av.VideoFrame(1080, 1920, format="yuv420p")

n = 1000
t_aligned = timeit.repeat(lambda: aligned.to_ndarray(), repeat=5, number=n)
t_padded = timeit.repeat(lambda: padded.to_ndarray(), repeat=5, number=n)
t_aligned_yuv = timeit.repeat(lambda: aligned_yuv.to_ndarray(format="rgb24"), repeat=5, number=n)
t_padded_yuv = timeit.repeat(lambda: padded_yuv.to_ndarray(format="rgb24"), repeat=5, number=n)

t_aligned = np.array(t_aligned) / n
t_padded = np.array(t_padded) / n
t_aligned_yuv = np.array(t_aligned_yuv) / n
t_padded_yuv = np.array(t_padded_yuv) / n

print(f"Aligned (1920x1080 rgb24):   {t_aligned.mean() * 1e6:.2f} ± {t_aligned.std() * 1e6:.2f} µs")
print(f"Padded  (1080x1920 rgb24):   {t_padded.mean() * 1e6:.2f} ± {t_padded.std() * 1e6:.2f} µs")
print(f"Aligned (1920x1080 yuv420p): {t_aligned_yuv.mean() * 1e6:.2f} ± {t_aligned_yuv.std() * 1e6:.2f} µs")
print(f"Padded  (1080x1920 yuv420p): {t_padded_yuv.mean() * 1e6:.2f} ± {t_padded_yuv.std() * 1e6:.2f} µs")